#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Authors: kevin

import rospy
import os
import math
import time
import numpy as np
import serial
import threading
from vision_msgs.msg import BoundingBox2D, Detection2DArray,ObjectHypothesisWithPose,Detection2D
from kyle_robot_toolbox.robot_arm.arm4dof_uservo import Arm4DoFUServo
from geometry_msgs.msg import Pose, Point, Quaternion, Twist
from rikibot_grasp_race.msg import ObjectInfo, ObjectInfoArray
from std_msgs.msg import UInt8, Float64
from nav import *



class xArmGrasp(object):
    def __init__(self):
        rospy.init_node("xArmsGrasp")
        r = rospy.Rate(5)
        self.PI = 3.1415

        self.cg_x = rospy.get_param("~cg_x", 25)
        self.cg_y = rospy.get_param("~cg_y", 60)
        self.cg_z = rospy.get_param("~cg_z", 200)
        self.obj_ok = rospy.get_param("~obj_ok", 10)

        self.grasp_z = rospy.get_param("~grasp_z", -205)
        self.fix_height = rospy.get_param("~fix_height", True)

        self.place_pose = rospy.get_param("~place_pose", [120, 0, 80])
        self.fix_grasp_mode = rospy.get_param("~fix_grasp_mode", True)

        self.grasp_pos = rospy.get_param("~grasp_pos")
        self.place_pos = rospy.get_param("~place_pos")
        self.end_pos = rospy.get_param("~end_pos")


        self.obj_cnt = 0
        self.nav_step = 0
        #self.grasp_status = False
        self.grasp_finish = False

        self.move_goal_lock = False
        self.place_goal_lock = False
        self.end_goal_lock = False

        self.grasp_lock = False
        self.fix_grasp_lock = False

        self.grasp_obj = None

        #dir_path = os.path.dirname(os.path.realpath(__file__))
        rospy.Subscriber("/object_info", ObjectInfoArray, self.ObjectCallback)
        #dir_path = dir_path.replace('rikibot_grasp_race/node', 'rikibot_grasp_race/')
        #dir_path += 'config/'
        dir_path = '/home/rikibot/catkin_ws/src/rikibot_project/rikibot_grasp_race/config/'
        self.arm = Arm4DoFUServo(config_folder=dir_path, is_init_pose=False)
        self.arm.gripper_open()
        self.arm.set_tool_pose(pose_name="capture_image")
        
        if self.fix_grasp_mode is False:
            rospy.loginfo("start navigation move to grasp pos")
            self.nav = Nav()

        while not rospy.is_shutdown():
            if self.fix_grasp_mode is False:
                if self.move_goal_lock is True and self.grasp_finish is False:
                    if self.grasp_lock is True :
                        rospy.loginfo("start grasp object")
                        x, y, z = self.CameratoGraspPose(self.grasp_obj)
                        self.GraspObject(x, y, z)
                else:
                    if self.end_goal_lock is False:
                        self.MoveNavGrasp()

                if self.place_goal_lock is True and self.end_goal_lock is False:
                    self.PlaceObject()
                    rospy.loginfo("end pos")
                    self.nav.SendGoal(self.NavPoseConvert(self.end_pos))
                    self.end_goal_lock = True
            else:
                if self.grasp_lock is True:
                    rospy.loginfo("start grasp object")
                    x, y, z = self.CameratoGraspPose(self.grasp_obj)
                    self.GraspObject(x, y, z)

            r.sleep()

    
    #获取目标坐标点
    def ObjectCallback(self, object_msg):
        self.object_info = object_msg
        if self.fix_grasp_mode is True or self.move_goal_lock is True and self.grasp_lock is False:
            self.obj_cnt = self.obj_cnt + 1

        if self.obj_cnt == self.obj_ok :
            if self.fix_grasp_mode is True or self.grasp_lock is False:
                self.grasp_obj = self.object_info.detections[-1]
                self.grasp_lock = True
        

    def MoveNavGrasp(self):
        if self.nav_step == 0:
            rospy.loginfo("first step")
            self.nav.SendGoal(self.NavPoseConvert(self.grasp_pos))
        if self.nav_step == 1 and self.grasp_finish is True:
            rospy.loginfo("second step")
            self.nav.SendGoal(self.NavPoseConvert(self.place_pos))

        state =self.nav.NavState()
        if state == 1:
            self.nav_step = self.nav_step + 1

            if self.nav_step == 1:
                self.move_goal_lock = True
                self.obj_cnt = 0
            if self.nav_step == 2:
                self.place_goal_lock = True


    #抓取流程
    def GraspObject(self, x, y, z):
        gx, gy, gz, gpitch = self.arm.get_tool_pose()
        self.arm.gripper_open()
        time.sleep(1)
        self.arm.set_tool_pose([x, y, z])
        time.sleep(1)
        self.arm.gripper_close()
        time.sleep(1)
        gx, gy, gz, gpitch = self.arm.get_tool_pose()
        time.sleep(1)
        if self.fix_grasp_mode is True:
            rospy.loginfo("move grasp current gripper pose: %f, %f, %f, %f", \
                    gx, gy, gz, np.degrees(gpitch))
            self.arm.set_tool_pose(self.place_pose)
            time.sleep(1)
            self.arm.gripper_open()
            time.sleep(1)
            #self.arm.home()
            self.arm.set_tool_pose(pose_name="capture_image")
            time.sleep(2)
            self.grasp_lock = False
            self.obj_cnt = 0
        else:
            self.arm.set_tool_pose([120, 0 , 80])
            self.grasp_lock = False
            self.grasp_finish = True

    def PlaceObject(self):
        self.arm.set_tool_pose(self.place_pose)
        time.sleep(1)
        self.arm.gripper_open()
        time.sleep(1)
        #self.arm.home()
        self.arm.set_tool_pose(pose_name="capture_image")
        time.sleep(2)
        self.grasp_lock = False
        self.obj_cnt = 0
        


    #摄像头与机械臂的坐标转换
    def CameratoGraspPose(self, obj_info):
        rospy.loginfo("current object xyz: %d, %d, %d", \
                obj_info.point_x, obj_info.point_y, obj_info.point_z)
        gx, gy, gz, gpitch = self.arm.get_tool_pose()
        rospy.loginfo("current gripper pose: %f, %f, %f, %f", \
                gx, gy, gz, np.degrees(gpitch))

        grasp_px = gx + self.cg_y - obj_info.point_y
        if self.fix_height is True:
            grasp_pz = self.grasp_z
        else:
            grasp_pz = gz + self.cg_z - obj_info.point_z
        grasp_py = self.cg_x - obj_info.point_x
        rospy.loginfo("grasp pose gx : %f , gy: %f, gz: %f", grasp_px, grasp_py,  grasp_pz)
        return grasp_px, grasp_py, grasp_pz


    def NavPoseConvert(self, pos):
        roll = 0
        pitch = 0
        yaw = np.deg2rad(pos[2])
        qx = np.sin(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) - np.cos(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
        qy = np.cos(roll/2) * np.sin(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.cos(pitch/2) * np.sin(yaw/2)
        qz = np.cos(roll/2) * np.cos(pitch/2) * np.sin(yaw/2) - np.sin(roll/2) * np.sin(pitch/2) * np.cos(yaw/2)
        qw = np.cos(roll/2) * np.cos(pitch/2) * np.cos(yaw/2) + np.sin(roll/2) * np.sin(pitch/2) * np.sin(yaw/2)
        #print(qx, qy, qz, qw)

        return Pose(Point(pos[0], pos[1], 0), Quaternion(qx, qy, qz, qw))


    def shutdown(self):
        # Release handle to the webcam
        rospy.logwarn("now will shutdown object_xyz_node ...")



if __name__ == '__main__':
    try:
        driver = xArmGrasp()
        rospy.spin()
    except Exception as e:
        print("Init xArm Grasp Failed...")
        raise e
